import torch
import matplotlib.pyplot as plt
import numpy as np

# 读取数据
data = torch.load("dataset_1dot25_continus_walk.pt", weights_only=False)

num_samples = len(data)
print("总样本数量:", num_samples)

# 只画前 36 个
num_to_plot = min(100, num_samples)

plt.figure(figsize=(12, 6))

for idx in range(num_to_plot):
    sample = data[idx+1000]
    tgt = sample["tgt"] if isinstance(sample, dict) else sample[1]
    tgt = np.array(tgt)

    # 如果 tgt 是 (T, 1)，取第 0 列
    if tgt.ndim == 2 and tgt.shape[1] == 1:
        tgt = tgt[:, 0]

    plt.plot(tgt, label=f"Sample {idx}", linewidth=1)

plt.title(f"tgt of First {num_to_plot} Samples")
plt.xlabel("Timestep")
plt.ylabel("tgt Value")
plt.grid(True, linestyle='--', alpha=0.6)
plt.legend(ncol=3, fontsize=8)
plt.tight_layout()
plt.show()
